#!/bin/bash

# AIQuant系统SSL证书配置脚本
# 用于生成自签名证书或配置Let's Encrypt证书

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 配置变量
DOMAIN="aiquant.example.com"
API_DOMAIN="api.aiquant.example.com"
MONITOR_DOMAIN="monitor.aiquant.example.com"
SSL_DIR="/etc/nginx/ssl"
CERT_DIR="/etc/letsencrypt/live"
EMAIL="admin@example.com"

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查是否为root用户
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "此脚本需要root权限运行"
        exit 1
    fi
}

# 创建SSL目录
create_ssl_directories() {
    log_info "创建SSL证书目录..."
    mkdir -p "$SSL_DIR"
    mkdir -p "/var/www/certbot"
    chmod 755 "$SSL_DIR"
    log_success "SSL目录创建完成"
}

# 安装依赖
install_dependencies() {
    log_info "安装必要的依赖包..."
    
    if command -v apt-get &> /dev/null; then
        apt-get update
        apt-get install -y openssl certbot python3-certbot-nginx
    elif command -v yum &> /dev/null; then
        yum update -y
        yum install -y openssl certbot python3-certbot-nginx
    elif command -v dnf &> /dev/null; then
        dnf update -y
        dnf install -y openssl certbot python3-certbot-nginx
    else
        log_error "不支持的包管理器"
        exit 1
    fi
    
    log_success "依赖安装完成"
}

# 生成自签名证书
generate_self_signed_cert() {
    local domain=$1
    local cert_file="$SSL_DIR/${domain}.crt"
    local key_file="$SSL_DIR/${domain}.key"
    
    log_info "为 $domain 生成自签名证书..."
    
    # 创建证书配置文件
    cat > "/tmp/${domain}.conf" << EOF
[req]
default_bits = 2048
prompt = no
default_md = sha256
distinguished_name = dn
req_extensions = v3_req

[dn]
C=CN
ST=Beijing
L=Beijing
O=AIQuant
OU=IT Department
CN=$domain

[v3_req]
basicConstraints = CA:FALSE
keyUsage = nonRepudiation, digitalSignature, keyEncipherment
subjectAltName = @alt_names

[alt_names]
DNS.1 = $domain
DNS.2 = *.$domain
DNS.3 = localhost
IP.1 = 127.0.0.1
EOF

    # 生成私钥和证书
    openssl req -new -x509 -days 365 -nodes \
        -out "$cert_file" \
        -keyout "$key_file" \
        -config "/tmp/${domain}.conf" \
        -extensions v3_req
    
    # 设置权限
    chmod 644 "$cert_file"
    chmod 600 "$key_file"
    
    # 清理临时文件
    rm "/tmp/${domain}.conf"
    
    log_success "自签名证书生成完成: $cert_file"
}

# 生成所有自签名证书
generate_all_self_signed() {
    log_info "生成所有域名的自签名证书..."
    
    generate_self_signed_cert "$DOMAIN"
    generate_self_signed_cert "$API_DOMAIN"
    generate_self_signed_cert "$MONITOR_DOMAIN"
    
    # 生成CA bundle (自签名情况下指向自己)
    cp "$SSL_DIR/$DOMAIN.crt" "$SSL_DIR/ca-bundle.crt"
    
    log_success "所有自签名证书生成完成"
}

# 获取Let's Encrypt证书
get_letsencrypt_cert() {
    local domain=$1
    
    log_info "为 $domain 获取Let's Encrypt证书..."
    
    # 使用webroot模式获取证书
    certbot certonly \
        --webroot \
        --webroot-path=/var/www/certbot \
        --email "$EMAIL" \
        --agree-tos \
        --no-eff-email \
        --force-renewal \
        -d "$domain"
    
    if [[ $? -eq 0 ]]; then
        # 创建符号链接到nginx SSL目录
        ln -sf "$CERT_DIR/$domain/fullchain.pem" "$SSL_DIR/$domain.crt"
        ln -sf "$CERT_DIR/$domain/privkey.pem" "$SSL_DIR/$domain.key"
        ln -sf "$CERT_DIR/$domain/chain.pem" "$SSL_DIR/ca-bundle.crt"
        
        log_success "Let's Encrypt证书获取成功: $domain"
    else
        log_error "Let's Encrypt证书获取失败: $domain"
        return 1
    fi
}

# 获取所有Let's Encrypt证书
get_all_letsencrypt() {
    log_info "获取所有域名的Let's Encrypt证书..."
    
    # 检查域名是否可达
    for domain in "$DOMAIN" "$API_DOMAIN" "$MONITOR_DOMAIN"; do
        if ! nslookup "$domain" &> /dev/null; then
            log_warning "域名 $domain 无法解析，跳过Let's Encrypt证书获取"
            continue
        fi
        
        get_letsencrypt_cert "$domain"
    done
    
    log_success "Let's Encrypt证书获取完成"
}

# 设置证书自动续期
setup_cert_renewal() {
    log_info "设置证书自动续期..."
    
    # 创建续期脚本
    cat > "/usr/local/bin/renew-certs.sh" << 'EOF'
#!/bin/bash

# AIQuant SSL证书续期脚本
LOG_FILE="/var/log/cert-renewal.log"

echo "$(date): 开始证书续期检查" >> "$LOG_FILE"

# 续期证书
certbot renew --quiet --post-hook "systemctl reload nginx" >> "$LOG_FILE" 2>&1

if [[ $? -eq 0 ]]; then
    echo "$(date): 证书续期检查完成" >> "$LOG_FILE"
else
    echo "$(date): 证书续期失败" >> "$LOG_FILE"
fi
EOF

    chmod +x "/usr/local/bin/renew-certs.sh"
    
    # 添加到crontab
    (crontab -l 2>/dev/null; echo "0 2 * * * /usr/local/bin/renew-certs.sh") | crontab -
    
    log_success "证书自动续期设置完成"
}

# 验证证书
verify_certificates() {
    log_info "验证SSL证书..."
    
    for domain in "$DOMAIN" "$API_DOMAIN" "$MONITOR_DOMAIN"; do
        cert_file="$SSL_DIR/${domain}.crt"
        key_file="$SSL_DIR/${domain}.key"
        
        if [[ -f "$cert_file" && -f "$key_file" ]]; then
            # 检查证书有效性
            if openssl x509 -in "$cert_file" -text -noout &> /dev/null; then
                # 检查证书和私钥是否匹配
                cert_hash=$(openssl x509 -noout -modulus -in "$cert_file" | openssl md5)
                key_hash=$(openssl rsa -noout -modulus -in "$key_file" | openssl md5)
                
                if [[ "$cert_hash" == "$key_hash" ]]; then
                    # 获取证书信息
                    expiry_date=$(openssl x509 -in "$cert_file" -noout -enddate | cut -d= -f2)
                    subject=$(openssl x509 -in "$cert_file" -noout -subject | cut -d= -f2-)
                    
                    log_success "证书验证成功: $domain"
                    log_info "  主题: $subject"
                    log_info "  过期时间: $expiry_date"
                else
                    log_error "证书和私钥不匹配: $domain"
                fi
            else
                log_error "证书文件无效: $cert_file"
            fi
        else
            log_warning "证书文件不存在: $domain"
        fi
    done
}

# 生成DH参数
generate_dhparam() {
    local dhparam_file="$SSL_DIR/dhparam.pem"
    
    if [[ ! -f "$dhparam_file" ]]; then
        log_info "生成DH参数文件 (这可能需要几分钟)..."
        openssl dhparam -out "$dhparam_file" 2048
        chmod 644 "$dhparam_file"
        log_success "DH参数文件生成完成"
    else
        log_info "DH参数文件已存在，跳过生成"
    fi
}

# 创建nginx SSL配置片段
create_nginx_ssl_config() {
    log_info "创建Nginx SSL配置片段..."
    
    cat > "/etc/nginx/snippets/ssl-params.conf" << EOF
# SSL配置参数
ssl_protocols TLSv1.2 TLSv1.3;
ssl_ciphers ECDHE-RSA-AES128-GCM-SHA256:ECDHE-RSA-AES256-GCM-SHA384:ECDHE-RSA-AES128-SHA256:ECDHE-RSA-AES256-SHA384;
ssl_prefer_server_ciphers off;
ssl_session_cache shared:SSL:10m;
ssl_session_timeout 10m;
ssl_session_tickets off;

# OCSP Stapling
ssl_stapling on;
ssl_stapling_verify on;
ssl_trusted_certificate $SSL_DIR/ca-bundle.crt;

# DH参数
ssl_dhparam $SSL_DIR/dhparam.pem;

# 安全头
add_header Strict-Transport-Security "max-age=63072000; includeSubDomains; preload" always;
add_header X-Frame-Options DENY always;
add_header X-Content-Type-Options nosniff always;
add_header X-XSS-Protection "1; mode=block" always;
add_header Referrer-Policy "strict-origin-when-cross-origin" always;
EOF

    log_success "Nginx SSL配置片段创建完成"
}

# 测试SSL配置
test_ssl_config() {
    log_info "测试SSL配置..."
    
    # 测试nginx配置
    if nginx -t &> /dev/null; then
        log_success "Nginx配置测试通过"
    else
        log_error "Nginx配置测试失败"
        nginx -t
        return 1
    fi
    
    # 重新加载nginx
    if systemctl reload nginx &> /dev/null; then
        log_success "Nginx重新加载成功"
    else
        log_error "Nginx重新加载失败"
        return 1
    fi
}

# 显示帮助信息
show_help() {
    cat << EOF
AIQuant SSL证书配置脚本

用法: $0 [选项]

选项:
    -h, --help              显示此帮助信息
    -d, --domain DOMAIN     设置主域名 (默认: $DOMAIN)
    -e, --email EMAIL       设置邮箱地址 (默认: $EMAIL)
    -s, --self-signed       生成自签名证书
    -l, --letsencrypt       获取Let's Encrypt证书
    -r, --renew             续期现有证书
    -v, --verify            验证现有证书
    --install-deps          仅安装依赖
    --setup-renewal         仅设置自动续期

示例:
    $0 --self-signed                    # 生成自签名证书
    $0 --letsencrypt -d example.com     # 获取Let's Encrypt证书
    $0 --verify                         # 验证现有证书
EOF
}

# 主函数
main() {
    local action=""
    
    # 解析命令行参数
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            -d|--domain)
                DOMAIN="$2"
                API_DOMAIN="api.$DOMAIN"
                MONITOR_DOMAIN="monitor.$DOMAIN"
                shift 2
                ;;
            -e|--email)
                EMAIL="$2"
                shift 2
                ;;
            -s|--self-signed)
                action="self-signed"
                shift
                ;;
            -l|--letsencrypt)
                action="letsencrypt"
                shift
                ;;
            -r|--renew)
                action="renew"
                shift
                ;;
            -v|--verify)
                action="verify"
                shift
                ;;
            --install-deps)
                action="install-deps"
                shift
                ;;
            --setup-renewal)
                action="setup-renewal"
                shift
                ;;
            *)
                log_error "未知选项: $1"
                show_help
                exit 1
                ;;
        esac
    done
    
    # 如果没有指定操作，显示帮助
    if [[ -z "$action" ]]; then
        show_help
        exit 1
    fi
    
    # 检查root权限
    check_root
    
    # 执行相应操作
    case $action in
        "install-deps")
            install_dependencies
            ;;
        "self-signed")
            install_dependencies
            create_ssl_directories
            generate_all_self_signed
            generate_dhparam
            create_nginx_ssl_config
            verify_certificates
            test_ssl_config
            ;;
        "letsencrypt")
            install_dependencies
            create_ssl_directories
            get_all_letsencrypt
            generate_dhparam
            create_nginx_ssl_config
            setup_cert_renewal
            verify_certificates
            test_ssl_config
            ;;
        "renew")
            certbot renew --post-hook "systemctl reload nginx"
            verify_certificates
            ;;
        "verify")
            verify_certificates
            ;;
        "setup-renewal")
            setup_cert_renewal
            ;;
    esac
    
    log_success "SSL证书配置完成！"
}

# 运行主函数
main "$@"